import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from cv2 import cv2
from tensorflow.keras.layers import Conv2D, BatchNormalization, Activation, MaxPool2D, Dropout, Flatten, Dense
from tensorflow.python.keras.models import Model

model_save_path = './checkpoint/DCF.ckpt'
class Baseline(Model):
    def __init__(self):
        super(Baseline, self).__init__()
        self.c1 = Conv2D(filters=6, kernel_size=(5, 5), padding='same')  # 卷积层
        self.b1 = BatchNormalization()  # BN层
        self.a1 = Activation('relu')  # 激活层
        self.p1 = MaxPool2D(pool_size=(2, 2), strides=2, padding='same')  # 池化层
        # self.d1 = Dropout(0.2)  # dropout层

        self.flatten = Flatten()
        self.f1 = Dense(128, activation='relu')
        # self.d2 = Dropout(0.2)
        self.f2 = Dense(10, activation='relu')
        self.f3 = Dense(3, activation='softmax')

    def call(self, x):
        x = self.c1(x)
        x = self.b1(x)
        x = self.a1(x)
        x = self.p1(x)
        # x = self.d1(x)

        x = self.flatten(x)
        x = self.f1(x)
        # x = self.d2(x)
        y = self.f2(x)
        y = self.f3(y)
        return y
model = Baseline()
model.load_weights(model_save_path)

# image_path = input("the path of test picture:")
image_path = 'flower.png'
image = cv2.imread(image_path)
image = cv2.resize(image, (28, 28))
plt.imshow(image)
img_arr = np.array(image, dtype=float)

x_predict = img_arr[tf.newaxis, ...]
result = model.predict(x_predict)
print(result)
pred = tf.argmax(result, axis=1).numpy()[0]
print('\n')
tf.print('Predict result: ', end='')
tf.print(['cat', 'dog', 'flower'][pred])
plt.pause(1)
plt.close()
